# 图像拼接 map1 map2
import cv2
import numpy as np


def get_homo(img1, img2):
    sift = cv2.SIFT_create()
    k1, p1 = sift.detectAndCompute(img1, None)
    k2, p2 = sift.detectAndCompute(img2, None)

    bf = cv2.BFMatcher()
    matches = bf.knnMatch(p1, p2, k=2)
    good = []
    for m1, m2 in matches:
        if m1.distance < 0.8 * m2.distance:
            good.append(m1)

    if len(good) > 8:
        img1_pts = []
        img2_pts = []
        for m in good:
            img1_pts.append(k1[m.queryIdx].pt)
            img2_pts.append(k2[m.trainIdx].pt)
        img1_pts = np.float32(img1_pts).reshape(-1, 1, 2)
        img2_pts = np.float32(img2_pts).reshape(-1, 1, 2)
        H, mask = cv2.findHomography(img1_pts, img2_pts, cv2.RANSAC, 5.0)
        return H
    else:
        print("points is not enough 8!")
        exit()


def stitch_img(img1, img2, H):
    h1, w1 = img1.shape[:2]
    h2, w2 = img1.shape[:2]
    img1_point = np.float32([[0, 0], [0, h1], [w1, h1], [w1, 0]]).reshape(-1, 1, 2)
    img2_point = np.float32([[0, 0], [0, h2], [w2, h2], [w2, 0]]).reshape(-1, 1, 2)
    img1_trans = cv2.perspectiveTransform(img1_point, H)
    print(img1_trans)
    result_point = np.concatenate((img2_point, img1_trans), axis=0)
    print(result_point)
    [x_min, y_min] = np.int32(result_point.min(axis=0).ravel() - 0.5)
    print([x_min, y_min])
    [x_max, y_max] = np.int32(result_point.max(axis=0).ravel() + 0.5)
    print([x_max, y_max])
    trans_dist = [-x_min, -y_min]
    trans_array = np.array([[1, 0, trans_dist[0]], [0, 1, trans_dist[1]], [0, 0, 1]])

    res_img = cv2.warpPerspective(img1, trans_array.dot(H), (x_max - x_min, y_max - y_min))
    res_img[trans_dist[1]:trans_dist[1] + h2, trans_dist[0]:trans_dist[0] + w2] = img2
    return res_img


def stitch_img2(img1, img2, H):
    h1, w1 = img1.shape[:2]
    h2, w2 = img1.shape[:2]
    res_img = cv2.warpPerspective(img1, H, (w1 + w2, h1))
    res_img[0:h2, 0:w2] = img2
    return res_img


img1 = cv2.imread('../asset/map1.jpg')
img2 = cv2.imread('../asset/map2.jpg')

input = np.hstack((img1, img2))
cv2.imshow('input', input)
H = get_homo(img1, img2)
print(H)
res_img = stitch_img2(img1, img2, H)
cv2.imshow('result', res_img)
cv2.waitKey(0)
